{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "from IPython import display\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 创建游戏环境\n",
    "def get_state(row, col):\n",
    "    if row != 3:\n",
    "        return \"ground\"\n",
    "    if row != 3 and col == 0:\n",
    "        return \"ground\"\n",
    "    if row == 3 and col == 11:\n",
    "        return \"terminal\"\n",
    "    if row == 3 and col == 0:\n",
    "        return \"ground\"\n",
    "    return \"trap\"\n",
    "\n",
    "\n",
    "def move(row, col, action):\n",
    "    if get_state(row, col) in [\"trap\", \"terminal\"]:\n",
    "        return row, col, 0\n",
    "\n",
    "    # 向上\n",
    "    if action == 0:\n",
    "        row -= 1\n",
    "\n",
    "    # 向下\n",
    "    if action == 1:\n",
    "        row += 1\n",
    "\n",
    "    # 向左\n",
    "    if action == 2:\n",
    "        col -= 1\n",
    "\n",
    "    # 向右\n",
    "    if action == 3:\n",
    "        col += 1\n",
    "\n",
    "    # 不允许走到地图外面\n",
    "    row = max(0, row)\n",
    "    row = min(3, row)\n",
    "    col = max(0, col)\n",
    "    col = min(11, col)\n",
    "\n",
    "    # 陷阱 奖励-100，否则是-1\n",
    "    reward = -1\n",
    "    if get_state(row, col) == \"trap\":\n",
    "        reward = -100\n",
    "\n",
    "    return row, col, reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4, 12, 4)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 初始化Q矩阵和3个list\n",
    "\n",
    "Q = np.zeros([4, 12, 4])\n",
    "\n",
    "# 初始化3个list，存储状态、动作、反馈的历史数据\n",
    "state_list = []\n",
    "action_list = []\n",
    "reward_list = []\n",
    "Q.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 动作函数 ——贪婪算法\n",
    "# 根据状态选择一个动作\n",
    "def get_action(row, col):\n",
    "    # 有一个小的概率随机选择动作\n",
    "    if random.random() < 0.1:\n",
    "        return random.choice(range(4))\n",
    "\n",
    "    # 否则选择分数最高的动作\n",
    "    return Q[row, col].argmax()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 获取5个时间步分别的分数\n",
    "def get_update_list(next_row, next_col, next_action):\n",
    "    # 初始化的target 是最后一个state和最后一个action的分数\n",
    "    target = Q[next_row, next_col, next_action]\n",
    "\n",
    "    # 计算每一步的target\n",
    "    # 每一步的target等于下一步的target*0.9，再加上本步的reward\n",
    "    # 时间从后向前回溯，越以前的会累加的信息越多\n",
    "    # 例如，走的五个状态记为 1 2 3 4 5\n",
    "    # 1.先计算状态5 要执行动作的分数tar5 = Q[状态5，动作5]\n",
    "    # 2.再计算状态4 动作的分数tar4 = tar5*0.9 + 状态4的奖励\n",
    "    # 3.再计算状态3 动作的分数tar3 = tar4*0.9 + 状态3的奖励\n",
    "    # ....\n",
    "    # reversed(range(5)) = [4,3,2,1,0]\n",
    "    target_list = []\n",
    "    for i in reversed(range(5)):\n",
    "        target = 0.9 * target + reward_list[i]\n",
    "        target_list.append(target)\n",
    "\n",
    "    # 把顺序顺过来 [tar0,tar1,tar2,tar3,tar4]\n",
    "    target_list = list(reversed(target_list))\n",
    "\n",
    "    # 计算每一步的value\n",
    "    value_list = []\n",
    "    for i in range(5):\n",
    "        row, col = state_list[i]\n",
    "        action = action_list[i]\n",
    "        value_list.append(Q[row, col, action])\n",
    "\n",
    "    # 计算每一步的更新量\n",
    "    update_list = []\n",
    "    for i in range(5):\n",
    "        update = 0.1 * (target_list[i] - value_list[i]) \n",
    "        update_list.append(update)\n",
    "    \n",
    "    return update_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train():\n",
    "    for epoch in range(3000):\n",
    "        # 初始化当前位置\n",
    "        row = random.choice(range(4))\n",
    "        col = 0\n",
    "\n",
    "        # 初始化第一个动作\n",
    "        action = get_action(row, col)\n",
    "\n",
    "        # 计算反馈的和,这个数字应该越来越小\n",
    "        reward_sum = 0\n",
    "\n",
    "        # 初始化3个列表\n",
    "        state_list.clear()\n",
    "        action_list.clear()\n",
    "        reward_list.clear()\n",
    "\n",
    "        # 循环直到到达终点或掉进陷阱\n",
    "        while get_state(row, col) not in [\"terminal\", \"trap\"]:\n",
    "            # 执行动作\n",
    "            next_row, next_col, reward = move(row, col, action)\n",
    "            reward_sum += reward\n",
    "\n",
    "            # 求新位置的动作\n",
    "            next_action = get_action(next_row, next_col)\n",
    "\n",
    "            # 记录历史数据\n",
    "            state_list.append([row, col])\n",
    "            action_list.append(action)\n",
    "            reward_list.append(reward)\n",
    "\n",
    "            # 积累到5步以后开始更新参数\n",
    "            if len(state_list) == 5:\n",
    "                #计算分数\n",
    "                update_list = get_update_list(next_row,next_col,next_action)\n",
    "\n",
    "                #只更新第一步的分数\n",
    "                row,col = state_list[0]\n",
    "                action = action_list[0]\n",
    "                update= update_list[0]\n",
    "\n",
    "                Q[row,col,action] += update\n",
    "\n",
    "                #移除第一步，这样在下次循环时保持列表是5个元素\n",
    "                state_list.pop(0)\n",
    "                action_list.pop(0)\n",
    "                reward_list.pop(0)\n",
    "\n",
    "            #更新当前位置\n",
    "            row = next_row\n",
    "            col = next_col\n",
    "            action = next_action\n",
    "\n",
    "            #走到终点以后，更新剩下的步数的update\n",
    "        for i in range(len(state_list)):\n",
    "                row,col = state_list[i]\n",
    "                action = action_list[i]\n",
    "                update= update_list[i]\n",
    "                Q[row,col,action] += update\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 打印游戏，方便测试\n",
    "def show(row, col, action):\n",
    "    graph = [\"⬜\"] * 37 + [\"❌\"] * 10 + [\"🚩\"]\n",
    "    action = {0: \"🔺\", 1: \"🔻\", 2: \"👈\", 3: \"👉\"}[action]\n",
    "    graph[row * 12 + col] = action\n",
    "    graph = \"\".join(graph)\n",
    "\n",
    "    for i in range(0, 4 * 12, 12):\n",
    "        print(graph[i : i + 12])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜\n",
      "⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜\n",
      "⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜⬜🔻\n",
      "⬜❌❌❌❌❌❌❌❌❌❌🚩\n"
     ]
    }
   ],
   "source": [
    "from IPython import display\n",
    "import time\n",
    "\n",
    "\n",
    "def test():\n",
    "    row = random.choice(range(4))\n",
    "    col = 0\n",
    "\n",
    "    # 至多玩N步\n",
    "    for _ in range(200):\n",
    "        # 获取当前状态，若当前状态在终点或者掉进陷阱则终止\n",
    "        if get_state(row, col) in [\"trap\", \"terminal\"]:\n",
    "            break\n",
    "\n",
    "        # 选择最优的动作\n",
    "        action = Q[row, col].argmax()\n",
    "\n",
    "        # 打印这个动作\n",
    "        display.clear_output(wait=True)\n",
    "        time.sleep(0.1)\n",
    "        show(row, col, action)\n",
    "\n",
    "        # 执行动作\n",
    "        row, col, reward = move(row, col, action)\n",
    "test()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "↓←→→→→→→→→→↓\n",
      "→→→↑←↑↑↑↑↑↑↓\n",
      "↑←↑→↑←↑↑↑↑→↓\n",
      "↑↑↑↑↑↑↑↑↑↑↑↑\n"
     ]
    }
   ],
   "source": [
    "# 打印策略\n",
    "##行，列，动作\n",
    "tactic = Q.argmax(axis=2)\n",
    "for row in tactic:\n",
    "    for element in row:\n",
    "        if element == 0:\n",
    "            print(\"↑\", end=\"\")\n",
    "        elif element == 1:\n",
    "            print(\"↓\", end=\"\")\n",
    "        elif element == 2:\n",
    "            print(\"←\", end=\"\")\n",
    "        elif element == 3:\n",
    "            print(\"→\", end=\"\")\n",
    "    print()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Gym",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
